package com.btm.btmchat.service;

import com.btm.btmchat.util.PromptUtils;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

/**
 * AI服务类，提供与AI聊天模型交互的功能
 */
@Service
public class AIService {
    private final OpenAiChatModel openAiChatModel;
    private final ChatClient chatClient;
    private final ChatClient chatClientAll;
    private final InMemoryChatMemory inMemoryChatMemory;
    private final PromptUtils promptUtils;

    /**
     * 构造函数，初始化AI聊天模型和聊天客户端
     * @param openAiChatModel OpenAI聊天模型实例
     * @param chatClientBuilder 聊天客户端构建器
     */
    public AIService(OpenAiChatModel openAiChatModel, ChatClient.Builder chatClientBuilder, PromptUtils promptUtils) {
        this.openAiChatModel = openAiChatModel;
        this.promptUtils = promptUtils;
        
        // 创建内存会话管理器
        this.inMemoryChatMemory = new InMemoryChatMemory();
        
        this.chatClient = chatClientBuilder
                // 注册可以使用的MCP工具
                .defaultTools(
                        // 依赖mcp-server-mysql服务
                        "executeQuerySql","searchTableStructure","queryOrganizationByName","searchBusinessTerm"
                        // 依赖operatmysql服务
//                        "execute_sql","get_chinese_initials","get_table_name","get_table_desc","get_table_index","get_lock_tables"
                )
                // 系统提示词
                .defaultSystem(promptUtils.getMysqlSystemPrompt())
                .defaultAdvisors(new MessageChatMemoryAdvisor(this.inMemoryChatMemory), new SimpleLoggerAdvisor())
                .build();



        this.chatClientAll = chatClientBuilder
                // 注册可以使用的MCP工具
                .defaultTools(
                        // 依赖mcp-server-mysql服务
                        "executeAllSql","searchTableStructure","queryOrganizationByName","searchBusinessTerm"
                        // 依赖operatmysql服务
//                        "execute_sql","get_chinese_initials","get_table_name","get_table_desc","get_table_index","get_lock_tables"
                )
                // 系统提示词
                .defaultSystem(promptUtils.getMysqlSystemPrompt())
                .defaultAdvisors(new MessageChatMemoryAdvisor(this.inMemoryChatMemory), new SimpleLoggerAdvisor())
                .build();
    }

    /**
     * 简单聊天方法，无预设的提示词，直接调用AI模型获取响应
     * @param userInput 用户输入内容
     * @return AI模型的响应结果
     */
    public String simpleChat(String userInput) {
        return this.openAiChatModel.call(userInput);
    }

    /**
     * 使用系统预设提示的聊天方法
     * @param message 用户消息
     * @return 格式化后的信息
     */
    public String chatWithSystemPrompt(String message, String sessionId) {
        return chatClient.prompt()
                .advisors(advisor -> advisor.param("chat_memory_conversation_id", sessionId)
                        .param("chat_memory_response_size", 100))
                .user(promptUtils.getUserPrompt(message))
                .call()
                .content();
    }

    /**
     * 使用系统预设提示的流式聊天方法，返回一个Flux流，用于处理流式响应
     * @param message 用户消息
     * @return 包含AI模型流式响应的Flux<String>
     */
    public Flux<String> chatStreamWithSystemPrompt(String message) {
        return chatClient.prompt()
                .user(promptUtils.getUserPrompt(message))
                .stream()
                .content();
    }
    
    /**
     * 支持多轮对话的流式聊天方法
     * @param sessionId 会话ID
     * @param message 用户消息
     * @return 包含AI模型流式响应的Flux<String>
     */
    public Flux<String> chatStreamWithHistory(String sessionId, String message) {
        // 使用MessageChatMemoryAdvisor自动管理会话历史
        return chatClient.prompt()
                .advisors(advisor -> advisor.param("chat_memory_conversation_id", sessionId)
                        .param("chat_memory_response_size", 100))
                // 添加用户消息
                .user(promptUtils.getUserPrompt(message))
                .stream()
                .content();
    }



    /**
     * 支持多轮对话的流式聊天方法，支持执行任何类型的SQL
     * @param sessionId 会话ID
     * @param message 用户消息
     * @return 包含AI模型流式响应的Flux<String>
     */
    public Flux<String> chatStreamAllWithHistory(String sessionId, String message) {
        // 使用MessageChatMemoryAdvisor自动管理会话历史
        return chatClientAll.prompt()
                .advisors(advisor -> advisor.param("chat_memory_conversation_id", sessionId)
                        .param("chat_memory_response_size", 100))
                // 添加用户消息
                .user(promptUtils.getUserPrompt(message))
                .stream()
                .content();
    }

    public void cleanChatHistory(String sessionId) {
        inMemoryChatMemory.clear(sessionId);
    }
}